{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c772ce31",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3/dist-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.4\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2025-05-03 23:27:08,525] [INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/bin/ld: cannot find -laio: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n",
      "/usr/bin/ld: cannot find -lcufile: No such file or directory\n",
      "collect2: error: ld returned 1 exit status\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f96e798b580>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from accelerate import init_empty_weights, load_checkpoint_and_dispatch\n",
    "from transformers import AutoTokenizer, Gemma3ForCausalLM, TextStreamer\n",
    "import safetensors\n",
    "\n",
    "torch.autograd.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aac1ec36",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('google/gemma-3-27b-it')\n",
    "streamer = TextStreamer(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9177cc9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "996d76ecc13349fc82d752c7a85588bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/12 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of Gemma3ForCausalLM were not initialized from the model checkpoint at google/gemma-3-27b-it and are newly initialized: ['lm_head.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = Gemma3ForCausalLM.from_pretrained(\n",
    "    'google/gemma-3-27b-it', torch_dtype = torch.bfloat16, tie_word_embeddings=False,\n",
    "    device_map=\"auto\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90a1405c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=7)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.lm_head.weight.data.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ba83dc08",
   "metadata": {},
   "outputs": [],
   "source": [
    "v = model.model.embed_tokens.weight.data.clone().to(model.lm_head.weight.data.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8f468a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.lm_head.weight.data = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "822d3622",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict = model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ecd232a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from safetensors import safe_open\n",
    "\n",
    "folder = 'lora-embedding-128-gemma3-27b-malaysian-8k/checkpoint-150'\n",
    "\n",
    "f = safe_open(f\"{folder}/adapter_model.safetensors\", framework=\"pt\", device='cpu')\n",
    "keys = f.keys()\n",
    "keys = sorted(list(set([k.split('.lora')[0] for k in keys if '.lora' in k])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9a9c14b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████| 436/436 [00:00<00:00, 724.14it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "for k in tqdm(keys):\n",
    "    k_ori = k.replace('base_model.model.', '') + '.weight'\n",
    "    if 'embed_tokens' in k:\n",
    "        post_A = '.lora_embedding_A'\n",
    "        post_B = '.lora_embedding_B'\n",
    "    else:\n",
    "        post_A = '.lora_A.weight'\n",
    "        post_B = '.lora_B.weight'\n",
    "    A = k + post_A\n",
    "    B = k + post_B\n",
    "    \n",
    "    W = state_dict[k_ori]\n",
    "    if 'embed_tokens' not in k:\n",
    "        W = W.t()\n",
    "        \n",
    "    A = f.get_tensor(A).to(W.device)\n",
    "    B = f.get_tensor(B).to(W.device)\n",
    "    with torch.no_grad():\n",
    "        W.addmm_(A.t(), B.t(), alpha = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "add81b8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<bos><start_of_turn>user\n",
      "butoh hg<end_of_turn>\n",
      "<start_of_turn>model\n",
      "Terima kasih atas soalan yang diajukan. Dalam konteks ini, saya cuba untuk memberikan penjelasan yang jelas dan terperinci mengenai perkara yang dibincangkan.\n",
      "\n",
      "Pertama sekali, kita perlu memahami apa yang dimaksudkan dengan \"butoh hg\" dalam bahasa Malaysia. Ia merujuk kepada perkataan \"butoh\" yang bermaksud tidak sedar diri, tidak tahu malu atau tidak mempunyai kesopanan. Dalam konteks ini, kita boleh membincangkan tentang kepentingan memiliki kesopanan dan adab dalam interaksi sosial, khususnya dalam komunikasi antara individu.\n",
      "\n",
      "Kesopanan dan adab adalah aspek penting dalam kehidupan sosial. Mereka membantu dalam mewujudkan persekitaran yang harmoni, saling menghormati, dan mengekalkan hubungan yang baik antara individu. Dalam komunikasi, kesopanan dan adab berperanan sebagai mekanisme yang memastikan bahawa perbualan berjalan lancar, tanpa menyinggung perasaan mana-mana pihak.\n",
      "\n",
      "Dalam kehidupan seharian, kita perlu memahami batas-batas kesopanan dan adab dalam berkomunikasi dengan orang lain. Contohnya, kita tidak sepatutnya mengkritik atau menyindir orang lain secara terbuka, terutama di hadapan ramai. Sebaliknya, kita perlu membincangkan perkara tersebut secara peribadi dan secara konstruktif, dengan tujuan untuk membantu memperbaiki masalah atau memperbaiki situasi yang sedang dihadapi.\n",
      "\n",
      "Selain itu, kita juga perlu menghormati perbezaan pendapat dan pandangan orang lain. Setiap individu mempunyai latar belakang, pengalaman, dan pandangan yang berbeza. Oleh itu, kita perlu memahami dan menghargai perbezaan tersebut, serta berusaha untuk memahami perspektif orang lain sebelum membuatkan sebarang kesimpulan atau komentar yang mungkin menyinggung perasaan mereka.\n",
      "\n",
      "Dalam interaksi sosial, kita perlu menjadi peka terhadap tanda-tanda atau isyarat yang ditunjukkan oleh orang lain. Contohnya, jika seseorang menunjukkan rasa tidak selesa atau tidak senang dengan perbualan, kita perlu mengambil inisiatif untuk mengubah topik atau memberi ruang kepada mereka untuk mengungkapkan perasaan mereka.\n",
      "\n",
      "Dalam keseluruhan, memiliki kesopanan dan adab dalam komunikasi antara individu adalah sangat penting. Mereka membantu dalam mewujudkan persekitaran yang harmoni, saling menghormati, dan mengekalkan hubungan yang baik antara individu. Dalam masa yang sama, kita juga perlu menjadi peka terhadap tanda-tanda atau isyarat yang ditunjukkan oleh orang lain, serta berusaha untuk memahami perspektif mereka sebelum membuatkan sebarang kesimpulan atau komentar.\n",
      "\n",
      "Semoga penjelasan\n"
     ]
    }
   ],
   "source": [
    "d = [\n",
    "    {'role': 'user', 'content': 'butoh hg'}\n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(d, return_tensors = 'pt').to('cuda')\n",
    "generate_kwargs = dict(\n",
    "    input_ids=inputs,\n",
    "    max_new_tokens=512,\n",
    "    top_p=0.95,\n",
    "    top_k=50,\n",
    "    temperature=0.9,\n",
    "    do_sample=True,\n",
    "    repetition_penalty=1.05,\n",
    "    streamer=streamer\n",
    ")\n",
    "generation_output = model.generate(**generate_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2870d1bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<bos><start_of_turn>user\n",
      "generate 10 introduction for a voice call accomodation agent in local malay\n",
      "\n",
      "1. please include filler words also, eg, uhm, erm, 'huh', 'ah'\n",
      "2. do not start with fillers\n",
      "3. always use 'encik' atau 'puan' atau 'cik', do not use 'u', 'awak'\n",
      "4. you can use any random company name\n",
      "5. make sure cover a lot of different use cases.\n",
      "6. you can also try to add asking some personal details.\n",
      "\n",
      "A good example such as,\n",
      "- Terima kasih kerana menghubungi sistem bantuan suara kami. Adakah anda perlukan bantuan untuk membuat tempahan bilik, menyemak butiran perjalanan, atau permintaan tambahan seperti tuala atau makanan?\n",
      "- Selamat datang ke perkhidmatan digital hotel kami. Saya boleh bantu anda menyemak masa daftar masuk, status bilik, atau permintaan khas semasa penginapan anda.\n",
      "\n",
      "return in JSON, list of strings<end_of_turn>\n",
      "<start_of_turn>model\n",
      "1. \"Selamat datang ke Holiday Inn, Puan. Saya di sini untuk membantu anda. Boleh saya tanya apa yang membawa anda menghubungi kami pada masa ini?\"\n",
      "\n",
      "2. \"Terima kasih atas panggilan ke The Ritz-Carlton, Encik. Apa yang saya boleh bantu anda hari ini? Apakah berkenaan dengan penginapan di hotel kami?\"\n",
      "\n",
      "3. \"Uhm... Assalamualaikum, Cik. Saya dari Shangri-La Hotel. Erm... Boleah saya tahu bagaimana saya boleh membantu anda sekarang?\"\n",
      "\n",
      "4. \"Selamat pagi, Puan! Nama saya dari Hyatt Regency. Adakah anda mempunyai sebarang soalan atau bantuan yang diperlukan berkaitan dengan hotel kami?\"\n",
      "\n",
      "5. \"Ah, hai! Saya dari Four Seasons Resort. Adakah anda sedang merancang untuk bercuti di resort kami? Atau mungkin anda mempunyai pertanyaan mengenai kemudahan kami?\"\n",
      "\n",
      "6. \"Hello! Dari Grand Hyatt. Jika anda mempunyai sebarang pertanyaan mengenai hotel kami, sila beritahu saya. Erm... adakah anda sudah membuat tempahan sebelum ini?\"\n",
      "\n",
      "7. \"Assalamualaikum Encik, saya dari Sofitel Kuala Lumpur Damansara. Bagaimana saya boleh membantu anda hari ini? Adakah anda ingin membuat tempahan bilik atau terdapat sesuatu yang ingin anda tahu tentang hotel kami?\"\n",
      "\n",
      "8. \"Selamat petang, Puan. Dari Mandarin Oriental, Kuala Lumpur. Adakah anda memerlukan bantuan dengan tempahan atau maklumat berkaitan hotel kami? Ah... Saya ada catat bahawa anda telah tinggal di sini sekali sebelum ini.\"\n",
      "\n",
      "9. \"Uhm... Hai Cik. Nama saya dari Hilton Kuala Lumpur. Saya berbesar hati untuk membantu anda. Adakah anda mempunyai sebarang keperluan khusus untuk penginapan anda?\"\n",
      "\n",
      "10. \"Salam! Saya dari Le Meridien Kuala Lumpur Sentral. Erm... Bagaimana saya boleh membantu anda pada masa ini, Puan? Adakah anda ada keperluan untuk bilik keluarga atau mungkin bilik VIP?<end_of_turn>\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"\n",
    "generate 10 introduction for a voice call accomodation agent in local malay\n",
    "\n",
    "1. please include filler words also, eg, uhm, erm, 'huh', 'ah'\n",
    "2. do not start with fillers\n",
    "3. always use 'encik' atau 'puan' atau 'cik', do not use 'u', 'awak'\n",
    "4. you can use any random company name\n",
    "5. make sure cover a lot of different use cases.\n",
    "6. you can also try to add asking some personal details.\n",
    "\n",
    "A good example such as,\n",
    "- Terima kasih kerana menghubungi sistem bantuan suara kami. Adakah anda perlukan bantuan untuk membuat tempahan bilik, menyemak butiran perjalanan, atau permintaan tambahan seperti tuala atau makanan?\n",
    "- Selamat datang ke perkhidmatan digital hotel kami. Saya boleh bantu anda menyemak masa daftar masuk, status bilik, atau permintaan khas semasa penginapan anda.\n",
    "\n",
    "return in JSON, list of strings\n",
    "\"\"\"\n",
    "d = [\n",
    "    {'role': 'user', 'content': prompt}\n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(d, return_tensors = 'pt').to('cuda')\n",
    "generate_kwargs = dict(\n",
    "    input_ids=inputs,\n",
    "    max_new_tokens=512,\n",
    "    top_p=0.95,\n",
    "    top_k=50,\n",
    "    temperature=0.9,\n",
    "    do_sample=True,\n",
    "    repetition_penalty=1.05,\n",
    "    streamer=streamer\n",
    ")\n",
    "generation_output = model.generate(**generate_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89eb1547",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
